#!/usr/bin/env python3
#coding:utf-8

__author__ = 'xmxoxo<xmxoxo@qq.com>'

# 以图搜图 API 接口

from faces_lib import *
from VecSearch_faiss import *

from sanic import Sanic, response, request
import logging
import traceback

from sanic.worker.manager import WorkerManager
WorkerManager.THRESHOLD = 600

# 全局变量
dats, vsearch, persons, model, detector, predictor = [], None, None, None, None, None,

#-----------------------------------------

def model_init(workpath):
    ''' 初始化模型与向量
    '''
    print('正在初始化模型与向量...')

    # 加载向量文件
    print('正在加载向量文件...')
    vector_file = os.path.join(workpath, 'vec.npy')
    vectors = np.load(vector_file)
    print('vectors:',vectors.shape)

    # 2023/3/16 数据对应有错，需要重新核对，原因是：文件名并不是唯一
    # 加载数据
    faces_file = os.path.join(workpath, 'faces.txt')
    txts = readtxt(faces_file)
    dats = [x.split('/')[-1] for x in txts.splitlines()]
    print('人脸总数:', len(dats))

    # 加载人名字典
    persons = [x.split('/')[2] for x in txts.splitlines()]

    # 创建搜索对象
    print('正在加载向量搜索器...')
    begin = time.time()

    # 使用faiss搜索器
    print('Faiss Search Engine...')
    dim = vectors.shape[1]
    vsearch = VecSearch(dim=dim, nlist=200)
    vsearch.add(vectors)
    vsearch.reindex()

    total_time = (time.time() - begin) * 1000
    print('加载用时:%.0f毫秒' % (total_time))

    # 加载人脸检测和识别模型
    model = load_vggface_model()
    detector, predictor = load_detection_model()

    return dats, vsearch, persons, model, detector, predictor

# -----------------------------------------

async def api_face(request):
    global dats, vsearch, persons, model, detector, predictor

    ret = {
        "code": -1,
        "msg": "Error",
        }
    try:
        s_time = time.time()

        # 提取参数
        jdata = request.json
        logging.debug('api_face request:%s'% str(jdata))
        image_b64 = jdata.get("image", "")

        if image_b64 == '':
            return response.json(ret)

        pic = base64toarray(image_b64)
        print('pic:', type(pic), pic.shape)
        print('detector:', type(detector))
        print('predictor:', type(predictor))

        faces, oimage = face_detector(pic, detector, predictor, showwin=0, saveface=0)
        logging.debug('api_face find faces:%s'% len(faces))
        if len(faces)==0:
            ret = {
                "code": 0,
                "msg": "OK",
                "faces": [],
                "oimage": ""
                }
        else:
            faces_b64 = [array2base64(face) for face in faces]
            oimage_b64 = array2base64(oimage)
            ret = {
                "code": 0,
                "msg": "OK",
                "faces": faces_b64,
                "oimage": oimage_b64
                }

        return response.json(ret)
    except Exception as e:
        print(e)
        print(traceback.format_exc())
        return response.json(ret)

async def api_search(request):
    global dats, vsearch, persons, model, detector, predictor
    ret = {
        "code": -1,
        "msg": "Error",
        }
    try:
        begin = time.time()

        # 提取参数
        jdata = request.json
        logging.debug('api_search request:%s'% str(jdata))
        faces_b64 = jdata.get("faces", [])
        topn = jdata.get("topn", 5)

        #-----------------------------------------

        # 计算待搜索图片的向量
        print('正在生成向量...')
        faces = [base64toarray(x) for x in faces_b64]
        samples = load_images(faces)
        query = get_model_scores(model, samples)

        # 搜索向量，注意可能是多个结果
        D, I = vsearch.search(query, top=topn)
        pics = D.shape[0]

        ret_images = []
        match_results = []
        for n in range(pics):
            print('-'*40)
            rets = zip(I[n],D[n])
            results, titles = [], []
            for i, d in rets:
                pname = dats[i]
                # 查找人名
                name = persons[i]
                # 显示图片
                # cv2.imshow("source face", faces[n])

                picname = './output0/%s/%s' % (name, pname)
                results.append(picname)
                title_text = '相似度:%.2f, 姓名：%s, PID:%s, 图像名:%s' % (d, name, i, picname)
                print(title_text)

                title = '%s(%.2f)'% (name, d)
                titles.append(title)

            # 显示合成的结果图
            # ret_img = show_result(faces[n], results, titles, showimg=0)
            # ret_images.append(array2base64(ret_img))
            # 返回匹配结果
            match_results.append([results, titles])

        # 计算用时
        total_time = (time.time() - begin) * 1000
        print('搜索用时:%.0f毫秒' % (total_time))

        ret = {
            "code": 200,
            "msg": "OK",
            "results": ret_images,
            "match_results": match_results,
            "utime": total_time
            }
        return response.json(ret)

    except Exception as e:
        print(e)
        return response.json(ret)

async def setup_app(app, loop):
    global dats, vsearch, persons, model, detector, predictor

    # 初始化加载数据，数据目录
    # workpath = './test_images/'
    # workpath = './test_images_500/'
    workpath = './test_images_all/'
    dats, vsearch, persons, model, detector, predictor = model_init(workpath)

def start_server(args):
    # app = Sanic(__name__)
    # app = Sanic('face_search_api')

    '''
    app.register_listener(setup_app, "before_server_start")

    # API接口
    print('正在初始化路由...')
    app.add_route(handler=api_face, uri="/face", methods={"POST"})
    app.add_route(handler=api_search, uri="/search", methods={"POST"})
    '''

    app.config.GRACEFUL_SHUTDOWN_TIMEOUT = 120.0
    app.run(host=args.host, port=args.port, debug=args.debug, workers=args.workers)

#-----------------------------------------

app = Sanic('face_search_api')
app.register_listener(setup_app, "before_server_start")

# API接口
print('正在初始化路由...')
app.add_route(handler=api_face, uri="/face", methods={"POST"})
app.add_route(handler=api_search, uri="/search", methods={"POST"})
'''
'''


if __name__ == '__main__':
    pass
    parser = argparse.ArgumentParser(description='API接口服务端')
    parser.add_argument('--host', type=str, default="0.0.0.0", help='host')
    parser.add_argument('--port', type=int, default=9505, help='port')
    parser.add_argument('--workers', type=int, default=1, help='workers')
    parser.add_argument('--debug', type=int, default=1, help='debug')

    args = parser.parse_args()
    start_server(args)

    '''
    print('正在启动...')
    app.config.GRACEFUL_SHUTDOWN_TIMEOUT = 120.0
    app.run(host=args.host, port=args.port, debug=args.debug, workers=args.workers)
    '''


